import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
import binascii
import base64

class FileEncryptorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("文件加密工具")
        self.root.geometry("600x400")
        self.root.resizable(True, True)
        
        # 创建主框架
        self.main_frame = ttk.Frame(root, padding="20")
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 密钥生成部分
        key_frame = ttk.LabelFrame(self.main_frame, text="密钥管理", padding=10)
        key_frame.pack(fill=tk.X, pady=5)
        
        ttk.Button(key_frame, text="生成新密钥", command=self.generate_key).grid(row=0, column=0, padx=5)
        ttk.Button(key_frame, text="保存密钥到文件", command=self.save_key).grid(row=0, column=1, padx=5)
        ttk.Button(key_frame, text="加载密钥文件", command=self.load_key).grid(row=0, column=2, padx=5)
        
        self.key_entry = ttk.Entry(key_frame, width=70)
        self.key_entry.grid(row=1, column=0, columnspan=3, sticky=tk.EW, pady=10)
        
        # 文件操作部分
        file_frame = ttk.LabelFrame(self.main_frame, text="文件操作", padding=10)
        file_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # 加密部分
        ttk.Label(file_frame, text="源文件:").grid(row=0, column=0, sticky=tk.W)
        self.source_entry = ttk.Entry(file_frame, width=50)
        self.source_entry.grid(row=0, column=1, sticky=tk.EW, padx=5)
        ttk.Button(file_frame, text="浏览...", command=self.browse_source).grid(row=0, column=2)
        
        ttk.Label(file_frame, text="输出文件:").grid(row=1, column=0, sticky=tk.W)
        self.output_entry = ttk.Entry(file_frame, width=50)
        self.output_entry.grid(row=1, column=1, sticky=tk.EW, padx=5)
        ttk.Button(file_frame, text="浏览...", command=self.browse_output).grid(row=1, column=2)
        
        action_frame = ttk.Frame(file_frame)
        action_frame.grid(row=2, column=0, columnspan=3, pady=10)
        
        ttk.Button(action_frame, text="加密文件", command=self.encrypt_file).pack(side=tk.LEFT, padx=5)
        ttk.Button(action_frame, text="解密文件", command=self.decrypt_file).pack(side=tk.LEFT, padx=5)
        
        # 状态栏
        self.status_var = tk.StringVar()
        self.status_bar = ttk.Label(root, textvariable=self.status_var, relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
        # 设置列权重
        file_frame.columnconfigure(1, weight=1)
        self.main_frame.columnconfigure(0, weight=1)
        
        # 初始化密钥
        self.key = None
        self.status_var.set("就绪")

    def generate_key(self):
        """生成新的AES密钥"""
        self.key = get_random_bytes(32)  # AES-256
        key_hex = binascii.hexlify(self.key).decode('utf-8')
        self.key_entry.delete(0, tk.END)
        self.key_entry.insert(0, key_hex)
        self.status_var.set("已生成新密钥")

    def save_key(self):
        """将密钥保存到文件"""
        if not self.key:
            messagebox.showerror("错误", "请先生成或加载密钥")
            return
        
        file_path = filedialog.asksaveasfilename(
            defaultextension=".key",
            filetypes=[("Key Files", "*.key"), ("All Files", "*.*")]
        )
        if not file_path:
            return
        
        try:
            with open(file_path, 'wb') as f:
                f.write(self.key)
            self.status_var.set(f"密钥已保存到: {file_path}")
        except Exception as e:
            messagebox.showerror("错误", f"保存密钥失败: {str(e)}")

    def load_key(self):
        """从文件加载密钥"""
        file_path = filedialog.askopenfilename(
            filetypes=[("Key Files", "*.key"), ("All Files", "*.*")]
        )
        if not file_path:
            return
        
        try:
            with open(file_path, 'rb') as f:
                self.key = f.read()
            
            if len(self.key) != 32:
                messagebox.showerror("错误", "无效的密钥长度，必须为32字节(AES-256)")
                self.key = None
                return
            
            key_hex = binascii.hexlify(self.key).decode('utf-8')
            self.key_entry.delete(0, tk.END)
            self.key_entry.insert(0, key_hex)
            self.status_var.set(f"已加载密钥: {file_path}")
        except Exception as e:
            messagebox.showerror("错误", f"加载密钥失败: {str(e)}")

    def browse_source(self):
        """浏览源文件"""
        file_path = filedialog.askopenfilename()
        if file_path:
            self.source_entry.delete(0, tk.END)
            self.source_entry.insert(0, file_path)

    def browse_output(self):
        """浏览输出文件"""
        file_path = filedialog.asksaveasfilename()
        if file_path:
            self.output_entry.delete(0, tk.END)
            self.output_entry.insert(0, file_path)

    def encrypt_file(self):
        """加密文件"""
        if not self.validate_key():
            return
        
        source = self.source_entry.get()
        output = self.output_entry.get()
        
        if not source or not output:
            messagebox.showerror("错误", "请选择源文件和输出文件")
            return
        
        try:
            iv = get_random_bytes(AES.block_size)
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
            
            with open(source, 'rb') as f:
                plaintext = f.read()
            
            # 添加PKCS#7填充并加密
            ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))
            
            # 将IV和密文一起写入文件
            with open(output, 'wb') as f:
                f.write(iv)
                f.write(ciphertext)
            
            self.status_var.set(f"文件加密成功: {output}")
        except Exception as e:
            messagebox.showerror("错误", f"加密失败: {str(e)}")

    def decrypt_file(self):
        """解密文件"""
        if not self.validate_key():
            return
        
        source = self.source_entry.get()
        output = self.output_entry.get()
        
        if not source or not output:
            messagebox.showerror("错误", "请选择源文件和输出文件")
            return
        
        try:
            with open(source, 'rb') as f:
                iv = f.read(AES.block_size)
                ciphertext = f.read()
            
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
            plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)
            
            with open(output, 'wb') as f:
                f.write(plaintext)
            
            self.status_var.set(f"文件解密成功: {output}")
        except Exception as e:
            messagebox.showerror("错误", f"解密失败: {str(e)}")

    def validate_key(self):
        """验证密钥是否有效"""
        if not self.key:
            # 尝试从输入框获取密钥
            key_str = self.key_entry.get().strip()
            if not key_str:
                messagebox.showerror("错误", "请先生成或输入密钥")
                return False
            
            try:
                # 尝试解析十六进制格式的密钥
                self.key = binascii.unhexlify(key_str)
                if len(self.key) != 32:
                    messagebox.showerror("错误", "无效的密钥长度，必须为32字节(AES-256)")
                    self.key = None
                    return False
            except:
                messagebox.showerror("错误", "无效的密钥格式，请使用十六进制格式")
                return False
        return True

if __name__ == "__main__":
    root = tk.Tk()
    app = FileEncryptorApp(root)
    root.mainloop()